# This module declares all the PyBLOG language specific features
# @var_dist
# @var etc.

# python imports

# local imports

# pyblog imports
import glob

class VariableDefinition:
  def __init__(self, fn, dep_fn, is_stochastic):
    """
    A VariableDefinition provides information about the function underlying a
    variable and other related functions.

    fn - the user visible name of the function defining the variable
    dep_fn - the function which computes the dependency statement
    is_stochastic - is the dependency stochastic or deterministic
    """
    self.fn = fn
    self.dep_fn = dep_fn
    self.is_stochastic = is_stochastic
    # store the number of arguments to the variable
    self.argcount = dep_fn.func_code.co_argcount
    # initially we assume that a variable is not a switching variable
    self.is_switching = False
    # initially we don't know whether instantiations of this variable can be
    # grouped or not
    self.groupby = None

def var_dist(func):
  """
  var_dist = variable with a distribution

  Declares the subsequent function to be a pyblog random variable with
  the return value being the distribution of the variable
  """
  # check that there are no default arguments
  if func.func_defaults is not None:
    raise Exception("PyBLOG variables can't have default arguments")
  
  # also, no *args or **kwargs
  if func.func_code.co_flags & (4 | 8):
    raise Exception("PyBLOG variables can't have list or dict arguments")

  # now, define a wrapper
  def wrapper(*args):
    if glob.curr_world is None:
      return RandomVariable(wrapper, *args)
    else:
      return glob.curr_world.query_var((wrapper,) + tuple(args))

  # safe decorator practice -- copy name, attribute and doc string from
  # original function also, the module name for the default help string
  wrapper.__name__ = func.__name__
  wrapper.__dict__ = func.__dict__
  wrapper.__doc__ = func.__doc__

  # store the definition of this variable
  glob.var_def[wrapper] = VariableDefinition(wrapper, func, True)
  
  return wrapper

def var(func):
  """
  Declares the subsequent function to be a pyblog random variable
  with a deterministic dependency
  """
  # check that there are no default arguments
  if func.func_defaults is not None:
    raise Exception("PyBLOG variables can't have default arguments")
  
  # also, no *args or **kwargs
  if func.func_code.co_flags & (4 | 8):
    raise Exception("PyBLOG variables can't have list or dict arguments")

  # now, define a wrapper
  def wrapper(*args):
    if glob.curr_world is None:
      return RandomVariable(wrapper, *args)
    else:
      return glob.curr_world.query_var((wrapper,) + tuple(args))

  # safe decorator practice -- copy name, attribute and doc string from
  # original function also, the module name for the default help string
  wrapper.__name__ = func.__name__
  wrapper.__dict__ = func.__dict__
  wrapper.__doc__ = func.__doc__
  
  # store the definition of this variable
  glob.var_def[wrapper] = VariableDefinition(wrapper, func, False)
  
  return wrapper

class RandomVariable:
  """
  A RandomVariable refers to a PyBLOG variable defined with @var_dist or
  @var. This refers to the name of the variable and not to its value in a
  PyBLOG world.

  A random variable == a value is short form for creating an Evidence
  object.
  """
  def __init__(self, fnwrap, *fnargs):
    # lookup the variable definition
    try:
      vardef = glob.var_def[fnwrap]
    except KeyError:
      raise TypeError("RandomVariable must be declated with @var_dist or @var")
    
    # check the number of arguments
    if len(fnargs) != vardef.argcount:
      raise ValueError("Function %s Expected %d arguments got %d" % \
                       (fnwrap.__name__, vardef.argcount, len(fnargs)))

    # record whether or not this variable has a stochastic dependency
    self.is_stochastic = vardef.is_stochastic
    
    # record the function name and its arguments
    self.fnwrap, self.fnargs = fnwrap, tuple(fnargs)

    # construct the varkey
    self.varkey = (self.fnwrap,) + self.fnargs

  def name(self):
    return self.varkey

  def fn_name(self):
    return self.varkey[0]

  def fn_args(self):
    return self.varkey[1:]
  
  def query(self):
    return self.fnwrap(*self.fnargs)
  
  def __str__(self):
    return self.fnwrap.__name__ + str(self.fnargs)

  # Using the == operator with a random variable generates an Evidence object
  
  # RandomVariable == val
  def __eq__(self, val):
    if val.__class__ == RandomVariable:
      raise TypeError("RandomVariables can't be equated to other random vars")
    return Evidence(self, val)

  # val == RandomVariable
  def __req__(self, val):
    if val.__class__ == RandomVariable:
      raise TypeError("RandomVariables can't be equated to other random vars")
    return Evidence(self, val)
  
class Evidence:
  def __init__(self, var, val):
    if var.__class__ != RandomVariable or val.__class__ == RandomVariable:
      raise TypeError("Evidence must be RandomVariable ==  Value")
    if not var.is_stochastic:
      raise TypeError("Evidence can only be given on stochastic random vars")
    self.var, self.val = var, val

  def __str__(self):
    return str(self.var)+"=="+str(self.val)

  def name(self):
    "Name of the random variable in the evidence"
    return self.var.varkey

  def value(self):
    "Value that the random variable is set to"
    return self.val

def gibbs_simple(variable):
  """
  A Gibbs simple proposer for a random variable gets the value of
  the random variable, its distribution and the values of all its children
  and it must return a new value for the variable. See the example below
  for the arguments of the proposer

  @var
  def z(i): return Bernoulli(.5)
  
  @gibbs_simple(z)
  def z_proposer(rand, z_dist, z_val, *z_children):
    ...
  
  """
  if variable not in glob.var_def or not glob.var_def[variable].is_stochastic:
    raise ("variable %s not declared with @var_dist" \
           % variable.__name__)
           
  def deco(gibbs_fn):
    # TODO: check the signature of the gibbs fn
    # store the gibbs sampler in the variable's definition
    glob.var_def[variable].gibbs_simple = gibbs_fn
    return gibbs_fn
    
  return deco


def gibbs_switching(variable):
  """
  A Gibbs switching proposer for a random variable gets the value of
  the random variable, its distribution and the values of all its children
  and it must return a new value for the variable. See the example below
  for the arguments of the proposer

  The main difference with Gibbs simple is that the structure of the world
  may change by changing a switching variable.

  @var
  def z(i): return Bernoulli(.5)
  
  @gibbs_switching(z)
  def z_proposer(rand, z_dist, z_val, *z_children):
    ...
  
  """
  # lookup the variable
  try:
    vardef = glob.var_def[variable]
  except KeyError:
    raise TypeError("Not a random variable %s" % variable.__name__)
  # verify that it has a stochastic dependency
  if not vardef.is_stochastic:
    raise TypeError("variable %s not declared with @var_dist" \
                    % variable.__name__)
  # mark it as a switching variable
  vardef.is_switching = True
  
  def deco(gibbs_fn):
    # TODO: check the signature of the gibbs fn
    # store the gibbs sampler in the variable's definition
    vardef.gibbs_switching = gibbs_fn
    return gibbs_fn
    
  return deco

class PyBLOGWorld:
  def __init__(self, var_values):
    self.var_values = var_values

  def __getitem__(self, var):
    if not isinstance(var, RandomVariable):
      raise TypeError("key must be a RandomVariable")
    return self.var_values[var.name()]

  def iteritems(self):
    for varkey, value in self.var_values.iteritems():
      yield RandomVariable(varkey[0], *varkey[1:]), value
    
@var
def pyblog_world():
  """
  Returns a hash table with all the variables and their values
  The hash key is of the form (fn. name, args).
  """
  return PyBLOGWorld(glob.curr_world.get_all_vars())
